package main

import (
	"fmt"
	"math/rand"
	"sync"
	"time"
)

// tee 函数将一个输入channel分成两个输出channel
// 每个输出channel都会收到来自输入channel的所有数据的副本
func tee(done <-chan struct{}, in <-chan interface{}) (_, _ <-chan interface{}) {
	out1 := make(chan interface{})
	out2 := make(chan interface{})

	go func() {
		defer close(out1)
		defer close(out2)

		for val := range orDone(done, in) {
			// 使用本地变量创建val的副本
			// 这样即使多个goroutine共享相同的内存地址，也能保证它们看到的是相同的值
			var out1, out2 = out1, out2
			var val = val

			// 将值发送给两个输出channel
			// 注意：我们需要处理以下两种情况
			// 1. 如果某个channel的接收者已停止工作，不应阻塞另一个channel
			// 2. 如果done通道被关闭，应该停止发送操作
			for i := 0; i < 2; i++ {
				select {
				case <-done:
					return
				case out1 <- val:
					// 第一个channel发送成功后，下次只发送给第二个
					out1 = nil
				case out2 <- val:
					// 第二个channel发送成功后，下次只发送给第一个
					out2 = nil
				}
			}
		}
	}()

	return out1, out2
}

// 更通用的teeN函数，支持任意数量的输出channel
func teeN(done <-chan struct{}, in <-chan interface{}, n int) []<-chan interface{} {
	if n <= 0 {
		return nil
	}

	// 创建n个输出channel
	outs := make([]chan interface{}, n)
	for i := range outs {
		outs[i] = make(chan interface{})
	}

	// 读取只读channel数组
	outReadOnly := make([]<-chan interface{}, n)
	for i := range outReadOnly {
		outReadOnly[i] = outs[i]
	}

	go func() {
		// 确保所有输出channel在函数退出时都被关闭
		defer func() {
			for _, ch := range outs {
				close(ch)
			}
		}()

		for val := range orDone(done, in) {
			// 创建值的局部副本
			val := val

			// 创建活跃channel的副本
			activeChans := append([]chan interface{}{}, outs...)

			// 尝试向所有活跃channel发送值
			for len(activeChans) > 0 {
				select {
				case <-done:
					return
				case activeChans[0] <- val:
					// 发送成功后从活跃列表中移除
					activeChans = activeChans[1:]
				}
			}
		}
	}()

	return outReadOnly
}

// orDone 函数接收一个done channel和一个输入channel
// 用于优雅地处理取消操作
func orDone(done <-chan struct{}, c <-chan interface{}) <-chan interface{} {
	valStream := make(chan interface{})

	go func() {
		defer close(valStream)

		for {
			select {
			case <-done:
				return
			case v, ok := <-c:
				if !ok {
					return
				}

				select {
				case valStream <- v:
				case <-done:
					return
				}
			}
		}
	}()

	return valStream
}

// 数据生成器，生成指定范围内的随机数
func dataGenerator(done <-chan struct{}, count int) <-chan interface{} {
	out := make(chan interface{})

	go func() {
		defer close(out)

		for i := 0; i < count; i++ {
			// 生成一个随机数据
			val := rand.Intn(100)

			select {
			case <-done:
				return
			case out <- val:
				// 模拟数据生成的时间间隔
				time.Sleep(100 * time.Millisecond)
			}
		}
	}()

	return out
}

// 数据处理器，对数据进行处理并跟踪ID
func processor(done <-chan struct{}, id int, in <-chan interface{}) <-chan interface{} {
	processed := make(chan interface{})

	go func() {
		defer close(processed)

		for val := range orDone(done, in) {
			// 处理数据：将数字转换为结构体，添加处理器ID信息
			if num, ok := val.(int); ok {
				result := struct {
					ProcessorID int
					Value       int
					Doubled     int
				}{
					ProcessorID: id,
					Value:       num,
					Doubled:     num * 2,
				}

				// 模拟处理时间
				procTime := time.Duration(rand.Intn(100)) * time.Millisecond
				time.Sleep(procTime)

				select {
				case <-done:
					return
				case processed <- result:
					fmt.Printf("处理器 %d: 处理了值 %d -> %d (耗时: %v)\n",
						id, num, result.Doubled, procTime)
				}
			}
		}
	}()

	return processed
}

// 合并多个channel的输出到一个channel
func fanIn(done <-chan struct{}, channels ...<-chan interface{}) <-chan interface{} {
	var wg sync.WaitGroup
	multiplexedStream := make(chan interface{})

	// 为每个输入channel启动一个goroutine
	multiplex := func(c <-chan interface{}) {
		defer wg.Done()
		for i := range orDone(done, c) {
			select {
			case <-done:
				return
			case multiplexedStream <- i:
			}
		}
	}

	// 计数要等待的goroutine数量
	wg.Add(len(channels))
	for _, c := range channels {
		go multiplex(c)
	}

	// 等待所有输入处理完毕后关闭输出
	go func() {
		wg.Wait()
		close(multiplexedStream)
	}()

	return multiplexedStream
}

// 示例场景：同一数据多种处理
// 比如日志同时发送到多个系统，或同一数据需要不同的处理方式
func processingSameDataMultipleWays() {
	done := make(chan struct{})
	defer close(done)

	// 创建数据源
	dataCh := dataGenerator(done, 5)

	// 将数据分成两路
	processingCh, loggingCh := tee(done, dataCh)

	// 数据处理
	processedCh := processor(done, 1, processingCh)

	// 日志记录（简单打印）
	go func() {
		for val := range orDone(done, loggingCh) {
			if num, ok := val.(int); ok {
				fmt.Printf("日志: 接收到值 %d\n", num)
			}
		}
	}()

	// 收集和显示处理结果
	for result := range processedCh {
		fmt.Printf("主程序: 收到处理结果 %v\n", result)
	}
}

// 示例场景：数据广播给多个相同类型的处理器
func broadcastToMultipleProcessors() {
	done := make(chan struct{})
	defer close(done)

	// 创建数据源
	dataCh := dataGenerator(done, 10)

	// 将数据广播给3个处理器
	processorChs := teeN(done, dataCh, 3)

	// 创建处理器
	processedChs := make([]<-chan interface{}, 3)
	for i := range processorChs {
		processedChs[i] = processor(done, i+1, processorChs[i])
	}

	// 合并所有处理器的输出
	mergedResults := fanIn(done, processedChs...)

	// 计数接收到的结果
	var resultCount int
	for result := range mergedResults {
		resultCount++
		fmt.Printf("主程序: 收到处理结果 #%d: %v\n", resultCount, result)
	}

	fmt.Printf("共处理了 %d 个结果\n", resultCount)
}

func main() {
	// 初始化随机种子
	rand.Seed(time.Now().UnixNano())

	fmt.Println("=== 示例1: 同一数据多种处理 ===")
	processingSameDataMultipleWays()

	fmt.Println("=== 示例2: 数据广播给多个相同类型的处理器 ===")
	broadcastToMultipleProcessors()

	// 演示通用tee方法的使用
	fmt.Println("=== 示例3: 使用teeN分配到N个输出 ===")
	done := make(chan struct{})
	dataCh := dataGenerator(done, 3)

	// 分发到4个输出
	outputs := teeN(done, dataCh, 4)

	// 为每个输出启动一个简单的消费者
	var wg sync.WaitGroup
	wg.Add(len(outputs))

	for i, ch := range outputs {
		go func(id int, c <-chan interface{}) {
			defer wg.Done()
			for val := range orDone(done, c) {
				fmt.Printf("消费者 %d: 收到值 %v\n", id, val)
				// 故意让消费者有不同的处理速度
				time.Sleep(time.Duration(50*id) * time.Millisecond)
			}
			fmt.Printf("消费者 %d: 完成处理\n", id)
		}(i+1, ch)
	}

	// 等待几秒后结束
	time.Sleep(3 * time.Second)
	fmt.Println("关闭done channel...")
	close(done)

	// 等待所有消费者完成
	wg.Wait()
	fmt.Println("所有消费者已完成")
}
